#!/usr/bin/env python
# -*- coding: utf-8 -*-

# import numpy as np
# import matplotlib.pyplot as plt
import re
import ast

OPERANDINFO_ATTRS = {
            'InList',
            'OutList',
            'Order',
            'Bitwidth',
            'AsmFormat',
}

p_operandinfo_attr = re.compile(r'^(?P<Attr>' + (r'|'.join(OPERANDINFO_ATTRS)) + r')\<(?P<Obj>[\w,\[\.\]]*)\>' + r'(?P<Value>=.*)?;$')

slist = [
    'InList<pg, ra, vb, pp>;',
    'InList<>;',
    'OutList<pu, pv>;',
    'OutList<>;',

    'Order<pg, pu, pv, ra, vb, pp> ;',
    'Order<pg,rd,PR,ra,vb> ;',
    'Order<pg, rd, M[ra, urb, offset], pp>;',
    'Order<pg, M[rs, soffset], M[ra, urb, offset], pp>;',
    'Order<pg, rd, C[vb, ra]>;',
    'Order<upg, UR[ura, uridx], urb>;',
    'Order<pg, rd, R[urb, ridx]>;',

    'Bitwidth<ra> = 32;',
    'Bitwidth<rd> = 32 + (num=="2")*32 + (num=="4")*96;',
    'Bitwidth<ra> = (ra.addrmode=="S32")*32 + (ra.addrmode=="U64")*64;',
    'Bitwidth<urb> = (urb.addrmode=="S32")*32 + (urb.addrmode=="U64")*64;',
    'Bitwidth<ra> =  ((mnk=="88128")*8 + (mnk=="168128")*16 + (mnk=="168256")*16) * ( (mnk=="88128")*128 + (mnk=="168128")*128 + (mnk=="168256")*256) / 32;',

    'AsmFormat<vb> = CvtFImm(vb, hfmt_v2);',
    'AsmFormat<vb.vsel> = CvtVSel(vb.vsel, itype);',
]

class OperandInfoParseError(Exception):
    pass

def parseList(tstr:str):

    t = ast.parse('[' + tstr + ']', mode='eval')
    tlist = []
    for elt in t.body.elts:
        if isinstance(elt, ast.Name):
            tlist.append(elt.id)
        elif isinstance(elt, ast.Subscript):
            sid = elt.value.id
            s_elts = [e.id for e in elt.slice.elts]
            tlist.append([sid, *s_elts])
        elif isinstance(elt, ast.Attribute):
            sid = elt.value.id
            sattr = elt.attr
            tlist.append(sid+'.'+sattr)
        else:
            raise OperandInfoParseError(f'Invalid ast type({type(elt)}) in operand list!')

    return tlist

astOpDict = {'Add' : '+', 'Mult' : '*', 'Div' : '/', 'Eq' : '=='}

def parseExpr(tstr:str):
    t = ast.parse(tstr, mode='eval')
    if not isinstance(t, ast.Expression):
        raise OperandInfoParseError(f'Invalid ast expr! {tstr}')

    # print(ast.dump(t, indent=2))
    vstack = []
    def push_v(v):
        if isinstance(v, ast.Name):
            vstack.append(v.id)
        elif isinstance(v, ast.Constant):
            vstack.append(v.value)
        elif isinstance(v, ast.BinOp):
            push_v(v.left)
            push_v(v.right)
            vop = v.op.__class__.__name__
            if vop not in astOpDict:
                raise OperandInfoParseError(f'Invalid BinOp({vop}) in operand info expr!')
            vstack.append(astOpDict[vop])
        elif isinstance(v, ast.Compare):
            push_v(v.left)
            push_v(v.comparators[0])
            cmp_op = v.ops[0].__class__.__name__
            if cmp_op not in astOpDict:
                raise OperandInfoParseError(f'Invalid CmpOp({cmp_op}) in operand info expr!')
            vstack.append(astOpDict[cmp_op])
        elif isinstance(v, ast.Attribute):
            vf = v.value.id
            vattr = v.attr
            vstack.append(vf + '.' + vattr)
        else:
            raise OperandInfoParseError(f'Invalid expr op {v}!')

    push_v(t.body)

    return vstack

def btest(slist):
    for s in slist:
        s = re.subn(r'\s', '', s)[0]
        print('#### ', s)
        res = p_operandinfo_attr.match(s)
        if res is not None:
            attr  = res.group('Attr')
            objs  = res.group('Obj')
            vals  = res.group('Value')
            if vals is None:
                vals = ''
            else:
                vals = vals.lstrip('=')

            print('  Attr : ', attr)
            print('  Objs : ', objs)
            print('  Vals : ', vals)

            if attr in {'InList', 'OutList', 'Order'}:
                print('Parsed obj: ')
                print(parseList(objs))
            elif attr in {'Bitwidth'}:
                print('Parsed val: ')
                print(parseExpr(vals))

            # if len(objs)>0:
            #     print('-'*80)
            #     print(f'Objs:  {objs}')
            #     print(ast.dump(ast.parse(objs), indent=2))
            # if len(vals)>0:
            #     print('-'*80)
            #     print(f'Vals: {vals}')
            #     print(ast.dump(ast.parse(vals), indent=2))
        else:
            print('None!!!!!!!!!!')
        print('\n')

def testEval():
    p = [ 32, "dtype", 5, "==", 32, "*", "+", "dtype", 6, "==", 96, "*", "+" ]

    d = {'dtype':7}
    vstack = []
    for v in p:
        if v == '==':
            r = vstack.pop()
            l = vstack.pop()
            vstack.append(l==r)
        elif v== '*':
            r = vstack.pop()
            l = vstack.pop()
            vstack.append(l*r)
        elif v== '+':
            r = vstack.pop()
            l = vstack.pop()
            vstack.append(l+r)
        elif v== '/':
            r = vstack.pop()
            l = vstack.pop()
            vstack.append(l//r)
        elif isinstance(v, str):
            vstack.append(d[v])
        else:
            vstack.append(v)

    print(vstack)


btest(slist)
